#include<string.h>
#include<stdlib.h>
#include<stdint.h>
#include<stdio.h>
#include<assert.h>
#include<fcntl.h>
#include "p2p-message.h"
#include "utils.h"
#include "p2p-file.h"
#include "p2p-peer.h"
#include "p2p-piece.h"
#include "p2p-config.h"
#include "my_log.h"

int handle_handshake_msg(int sockfd, const char *recv_buf, int bytes_recvd)
{
    uint8_t pstrlen;
    char *pstr;
    
    pstrlen = recv_buf[0];

    pstr = (char *)malloc(pstrlen * sizeof(char));
    memcpy(pstr, recv_buf + PSTRLEN_LEN, pstrlen);
    if(memcmp(pstr, "P2P", pstrlen)) {
        printf("Unknown Protocol. Dropping handshake packet\n");
        free(pstr);
        return 0;
    } 
    LOG(LOG_INFO, "Handshake message received from Peer:%d", sockfd);
    free(pstr);
    return 1;
}

int handle_request_msg(int sockfd, const char *buf, block_t *block) 
{
    block->index = my_ntohl(buf);
    block->begin = my_ntohl(buf + INDEX_LEN);
    block->length = my_ntohl(buf + INDEX_LEN + BEGIN_LEN);

    LOG(LOG_INFO, "Request message: Index:%d Begin:%d Length:%d", block->index, block->begin, block->length);
    return 0;
}

int handle_piece_msg(int sockfd, const char *buf, block_t *block, int len) 
{
    int bw, br = 0;
    
    block->index = my_ntohl(buf);
    block->begin = my_ntohl(buf + (br += INDEX_LEN));
    block->length = len - ID_LEN - INDEX_LEN - BEGIN_LEN;

    LOG(LOG_INFO, "Piece message: Index:%d Begin:%d Length:%d", block->index, block->begin, block->length);
    
    bw = write_file(FILE_, block->index * PIECE_SIZE + block->begin, 
                            buf + (br += BEGIN_LEN), block->length, O_WRONLY);

    if (bw == -1) {
        printf("Error writing to file: %s\n", FILE_);
        return -1;
    } 

    return 0;
}

int handle_bitfield_msg(int sockfd, const char *recv_buf, peer_t *peer, const uint8_t *piece_map, int no_of_pieces)
{
    uint32_t i, len;
    int count;

    len = my_ntohl(recv_buf);
    if (no_of_pieces != len - ID_LEN) {
        LOG(LOG_ERR, "Bitfield length not equal to # of pieces\n");
        return -1;
    }

    peer->piece_map = (uint8_t *)malloc(no_of_pieces * sizeof(uint8_t));
    memcpy(peer->piece_map, recv_buf + LEN_LEN + ID_LEN, no_of_pieces);

    // Construct piece list for this peer
    for (i = 0; i < no_of_pieces; i++) {
        // If the peer contains the piece
        if (recv_buf[LEN_LEN + ID_LEN + i]) {
            // If the piece is not present with client
            if (!piece_map[i]) {
                LOG(LOG_DBG, "Adding piece #%d to peer: %d\n", i, peer->psfd);
                // If this is the first piece in list
                if (!peer->piece_count) {
                    peer->piece_head = init_piece(i);
                    peer->piece_tail = peer->piece_head;
                } else {
                    peer->piece_tail = add_piece(peer->piece_tail, i);
                }
                peer->piece_count++;
            }
        }
    }

    return 0;
}

int handle_have_msg(int sockfd, const char *recv_buf, peer_t *peer, const uint8_t *piece_map)
{
    uint32_t index;

    index = my_ntohl(recv_buf);
    peer->piece_map[index] = 1;

    LOG(LOG_INFO, "Have message: Index:%d", index);
    // Add piece to this peer's piece list if the piece is not present with client
    if (!piece_map[index]) {
        LOG(LOG_DBG, "Adding piece #%d to peer: %d\n", index, peer->psfd);
        // If this is the first piece in list
        if (!peer->piece_count) {
            peer->piece_head = init_piece(index);
            peer->piece_tail = peer->piece_head;
        } else {
            peer->piece_tail = add_piece(peer->piece_tail, index);
        }
        peer->piece_count++;
    }

    return 0;
}

int send_handshake(int sockfd)
{
    uint8_t pstrlen = PSTRLEN;
    char *pstr = STR(PSTR);
    char *send_buf;
    // Message length = Length of PSTR + Length of PSTRLEN
    int i, msg_len = pstrlen + PSTRLEN_LEN;

    assert(strlen(pstr) == PSTRLEN);

    send_buf = (char *)malloc(msg_len * sizeof(char));
    memcpy(send_buf, &pstrlen, PSTRLEN_LEN);
    memcpy(send_buf + PSTRLEN_LEN, pstr, PSTRLEN);
   
    sendall(sockfd, send_buf, &msg_len);
    LOG(LOG_INFO, "Handshake message sent");
    
    free(send_buf);

    return 0;
}

int send_bitfield_msg(int sockfd)
{
    char *send_buf, *buffer;
    int msg_len, br;
    uint32_t len;
    uint8_t id;
    
    if ((br = read_file(META_FILE, 0, &buffer, -1)) == -1) {
        printf("[p2p-engine]: Meta file: %s not found\n", META_FILE);
        return -1;
    }
    
    msg_len = LEN_LEN + ID_LEN + br;
    send_buf = (char *)malloc(msg_len * sizeof(char));
    
    len = htonl(ID_LEN + br);
    id = BITFIELD_MSG;
    
    memcpy(send_buf, &len, LEN_LEN);
    memcpy(send_buf + LEN_LEN, &id, ID_LEN);
    memcpy(send_buf + LEN_LEN + ID_LEN, buffer, br);
   
    sendall(sockfd, send_buf, &msg_len);
    LOG(LOG_INFO, "Bitfield message sent");
    
    free(send_buf);
    free(buffer);
    
    return 0;
}

int send_piece_msg(int sockfd, const block_t* block)
{
    int br, msg_len;
    uint32_t len, index, begin;
    uint8_t id;

    char *send_buf, *file_buf;

    // Read the block requested from file
    br = read_file(FILE_, block->index * PIECE_SIZE + block->begin, &file_buf, block->length);

    if (br == -1) {
        LOG(LOG_ERR, "Error reading file: %s\n", FILE_);
        return -1;
    } else if (br < block->length) {
        LOG(LOG_ERR, "br not matching length requested");
        return -1;
    }

    // Construct the Piece message and send it
    msg_len = LEN_LEN + PIECE_MSG_LEN + block->length;
    send_buf = (char *)malloc(msg_len * sizeof(char));
    len = htonl(PIECE_MSG_LEN + block->length);
    id = PIECE_MSG;
    index = htonl(block->index);
    begin = htonl(block->begin);
    
    memcpy(send_buf, &len, LEN_LEN);
    memcpy(send_buf + LEN_LEN, &id, ID_LEN);
    memcpy(send_buf + LEN_LEN + ID_LEN, &index, INDEX_LEN);
    memcpy(send_buf + LEN_LEN + ID_LEN + INDEX_LEN, &begin, BEGIN_LEN);
    memcpy(send_buf + LEN_LEN + ID_LEN + INDEX_LEN + BEGIN_LEN, file_buf, block->length);
    
    sendall(sockfd, send_buf, &msg_len);
    LOG(LOG_INFO, "Piece message sent to peer:%d. Index: %d, Begin: %d, Length: %d", sockfd, block->index, block->begin, block->length);
    
    free(file_buf);
    free(send_buf);
    return 0;
}

int send_request_msg(int sockfd, const block_t *block)
{
    char *send_buf;
    int msg_len, bw = 0;
    uint32_t len, index, begin, length;
    uint8_t id;

    msg_len = LEN_LEN + REQUEST_MSG_LEN;
    send_buf = (char *)malloc(msg_len * sizeof(char));
    
    len = htonl(REQUEST_MSG_LEN);
    id = REQUEST_MSG;
    index = htonl(block->index);
    begin = htonl(block->begin);
    length = htonl(block->length);
    
    memcpy(send_buf + bw, &len, LEN_LEN);
    memcpy(send_buf + (bw += LEN_LEN), &id, ID_LEN);
    memcpy(send_buf + (bw += ID_LEN), &index, INDEX_LEN);
    memcpy(send_buf + (bw += INDEX_LEN), &begin, BEGIN_LEN);
    memcpy(send_buf + (bw += BEGIN_LEN), &length, LENGTH_LEN);
   
    sendall(sockfd, send_buf, &msg_len);
    LOG(LOG_INFO, "Request message sent to peer:%d. Index: %d, Begin: %d, Length: %d", sockfd, block->index, block->begin, block->length);
    
    free(send_buf);
    
    return 0;
}

int broadcast_have_msg(peer_t *peer, int index_r)
{
    char *send_buf;
    int msg_len, bw = 0;
    uint32_t len, index;
    uint8_t id;

    msg_len = LEN_LEN + HAVE_MSG_LEN;
    send_buf = (char *)malloc(msg_len * sizeof(char));
    
    len = htonl(HAVE_MSG_LEN);
    id = HAVE_MSG;
    index = htonl(index_r);
    
    memcpy(send_buf + bw, &len, LEN_LEN);
    memcpy(send_buf + (bw += LEN_LEN), &id, ID_LEN);
    memcpy(send_buf + (bw += ID_LEN), &index, INDEX_LEN);
   
    while (peer != NULL) {
        if (peer->piece_map != NULL) {
            if (!peer->piece_map[index_r]) {
                sendall(peer->psfd, send_buf, &msg_len);
                LOG(LOG_INFO, "Have message sent to peer:%d. Index: %d", peer->psfd, index_r);
            }
        }
        peer = peer->next;
    }
    
    free(send_buf);
    return 0;
}
